package cn.darksoul3.spark.cases

import org.apache.spark.Partitioner

import scala.collection.mutable

class TriplePartitioner(val titleTypes: Array[String]) extends Partitioner {

  val map = new mutable.HashMap[String, Int]()
  var i = 0
  for (elem <- titleTypes) {
    map(elem) = i
    i += 1
  }

  override def numPartitions: Int = titleTypes.length

  override def getPartition(key: Any): Int = {
    val tuple = key.asInstanceOf[(String, Int, Int)]
    val mapKey = tuple._1
    map(mapKey)
  }
}
